================================================================================
POPULATION WEIGHTING SUMMARY
================================================================================

Largest differences (weighted vs unweighted):

================================================================================
PHI COMPARISON: Weighted vs Unweighted Models
================================================================================

✅ Key Finding: Mean difference <0.002 indicates minimal impact of IPW on signature structure

================================================================================
SAMPLE DISEASE PHI COMPARISON
================================================================================

✅ Correlation between unweighted and weighted phi: 0.999948
   This high correlation confirms minimal impact of IPW on signature structure

import torch
import numpy as np
import matplotlib.pyplot as plt
import sys
import pandas as pd
from pathlib import Path

# Add path for utils
sys.path.append('/Users/sarahurbut/aladynoulli2/pyScripts')
from utils import calculate_pi_pred, softmax_by_k

print("="*80)
print("LAMBDA COMPARISON: Weighted vs Unweighted Models (Individual Level)")
print("="*80)

# Load weighted model (use first batch as example)
weighted_model_dir = Path("/Users/sarahurbut/Library/CloudStorage/Dropbox-Personal/batch_models_weighted_vec_censoredE/")
weighted_model_path = weighted_model_dir / "batch_00_model.pt"

# Load unweighted model (first batch)
unweighted_model_dir = Path("/Users/sarahurbut/Library/CloudStorage/Dropbox/censor_e_batchrun_vectorized/")
unweighted_model_path = unweighted_model_dir / "enrollment_model_W0.0001_batch_0_10000.pt"

if weighted_model_path.exists() and unweighted_model_path.exists():
    # Load models
    print(f"\nLoading weighted model: {weighted_model_path.name}")
    weighted_ckpt = torch.load(weighted_model_path, weights_only=False, map_location='cpu')
    
    print(f"Loading unweighted model: {unweighted_model_path.name}")
    unweighted_ckpt = torch.load(unweighted_model_path, weights_only=False, map_location='cpu')
    
    # Extract lambda (patient-specific signature loadings)
    # Lambda shape: [N, K, T]
    if 'model_state_dict' in weighted_ckpt:
        weighted_lambda = weighted_ckpt['model_state_dict']['lambda_'].detach()
        weighted_phi = weighted_ckpt['model_state_dict']['phi'].detach()
        weighted_kappa = weighted_ckpt['model_state_dict'].get('kappa', torch.tensor(1.0))
        if torch.is_tensor(weighted_kappa):
            weighted_kappa = weighted_kappa.item() if weighted_kappa.numel() == 1 else weighted_kappa.mean().item()
    else:
        weighted_lambda = weighted_ckpt['lambda_'].detach()
        weighted_phi = weighted_ckpt['phi'].detach()
        weighted_kappa = weighted_ckpt.get('kappa', torch.tensor(1.0))
        if torch.is_tensor(weighted_kappa):
            weighted_kappa = weighted_kappa.item() if weighted_kappa.numel() == 1 else weighted_kappa.mean().item()
    
    if 'model_state_dict' in unweighted_ckpt:
        unweighted_lambda = unweighted_ckpt['model_state_dict']['lambda_'].detach()
        unweighted_phi = unweighted_ckpt['model_state_dict']['phi'].detach()
        unweighted_kappa = unweighted_ckpt['model_state_dict'].get('kappa', torch.tensor(1.0))
        if torch.is_tensor(unweighted_kappa):
            unweighted_kappa = unweighted_kappa.item() if unweighted_kappa.numel() == 1 else unweighted_kappa.mean().item()
    else:
        unweighted_lambda = unweighted_ckpt['lambda_'].detach()
        unweighted_phi = unweighted_ckpt['phi'].detach()
        unweighted_kappa = unweighted_ckpt.get('kappa', torch.tensor(1.0))
        if torch.is_tensor(unweighted_kappa):
            unweighted_kappa = unweighted_kappa.item() if unweighted_kappa.numel() == 1 else unweighted_kappa.mean().item()
    
    print(f"\nWeighted lambda shape: {weighted_lambda.shape}")
    print(f"Unweighted lambda shape: {unweighted_lambda.shape}")
    
    # Ensure same shape (in case batch sizes differ)
    min_N = min(weighted_lambda.shape[0], unweighted_lambda.shape[0])
    weighted_lambda = weighted_lambda[:min_N]
    unweighted_lambda = unweighted_lambda[:min_N]
    
    print(f"Using {min_N} patients for comparison")
    
    # Compute correlation on ALL individual values (N×K×T)
    weighted_flat = weighted_lambda.numpy().flatten()
    unweighted_flat = unweighted_lambda.numpy().flatten()
    
    individual_correlation = np.corrcoef(weighted_flat, unweighted_flat)[0, 1]
    individual_mean_diff = np.abs(weighted_flat - unweighted_flat).mean()
    individual_max_diff = np.abs(weighted_flat - unweighted_lambda.numpy().flatten()).max()
    
    print(f"\nIndividual Lambda Comparison (N×K×T):")
    print(f"  Correlation: {individual_correlation:.6f}")
    print(f"  Mean absolute difference: {individual_mean_diff:.6f}")
    print(f"  Max absolute difference: {individual_max_diff:.6f}")
    
    # Also compute average lambda for heatmap visualization
    weighted_lambda_avg = weighted_lambda.mean(dim=0)  # [K, T]
    unweighted_lambda_avg = unweighted_lambda.mean(dim=0)  # [K, T]
    lambda_diff_avg = weighted_lambda_avg - unweighted_lambda_avg
    
    # Compute variance across patients for each signature×time
    weighted_lambda_var = weighted_lambda.var(dim=0)  # [K, T]
    unweighted_lambda_var = unweighted_lambda.var(dim=0)  # [K, T]
    
    # Plot comparison
    fig = plt.figure(figsize=(16, 12))
    
    # 1. Scatter plot: ALL individual values (N×K×T)
    ax1 = plt.subplot(2, 3, 1)
    # Subsample for visualization (too many points)
    n_sample = min(50000, len(weighted_flat))
    sample_idx = np.random.choice(len(weighted_flat), n_sample, replace=False)
    ax1.scatter(unweighted_flat[sample_idx], weighted_flat[sample_idx], alpha=0.1, s=0.5)
    ax1.plot([unweighted_flat.min(), unweighted_flat.max()], 
             [unweighted_flat.min(), unweighted_flat.max()], 'r--', alpha=0.7, linewidth=2)
    ax1.set_xlabel('Unweighted Lambda (Individual)', fontsize=11)
    ax1.set_ylabel('Weighted (IPW) Lambda (Individual)', fontsize=11)
    ax1.set_title(f'Individual Lambda: All N×K×T Values\nCorrelation: {individual_correlation:.4f}\n(n={n_sample:,} sampled)', 
                 fontsize=12, fontweight='bold')
    ax1.grid(True, alpha=0.3)
    
    # 2. Distribution of individual differences
    ax2 = plt.subplot(2, 3, 2)
    diff_flat = weighted_flat - unweighted_flat
    ax2.hist(diff_flat, bins=100, alpha=0.7, edgecolor='black')
    ax2.axvline(0, color='r', linestyle='--', linewidth=2, label='No difference')
    ax2.set_xlabel('Lambda Difference (Weighted - Unweighted)', fontsize=11)
    ax2.set_ylabel('Frequency', fontsize=11)
    ax2.set_title(f'Distribution of Individual Lambda Differences\nMean: {individual_mean_diff:.6f}', 
                 fontsize=12, fontweight='bold')
    ax2.legend()
    ax2.grid(True, alpha=0.3)
    
    # 3. Heatmap: Mean difference by signature and time
    ax3 = plt.subplot(2, 3, 3)
    im3 = ax3.imshow(lambda_diff_avg.numpy(), aspect='auto', cmap='RdBu_r', 
                     vmin=-lambda_diff_avg.abs().max().item(), 
                     vmax=lambda_diff_avg.abs().max().item())
    ax3.set_xlabel('Time (Age)', fontsize=11)
    ax3.set_ylabel('Signature', fontsize=11)
    ax3.set_title('Mean Lambda Difference: Weighted - Unweighted\n(Averaged across patients)', 
                 fontsize=12, fontweight='bold')
    plt.colorbar(im3, ax=ax3, label='Difference')
    
    # 4. Heatmap: Variance difference (shows where individual differences are largest)
    ax4 = plt.subplot(2, 3, 4)
    var_diff = weighted_lambda_var - unweighted_lambda_var
    im4 = ax4.imshow(var_diff.numpy(), aspect='auto', cmap='RdBu_r',
                     vmin=-var_diff.abs().max().item(),
                     vmax=var_diff.abs().max().item())
    ax4.set_xlabel('Time (Age)', fontsize=11)
    ax4.set_ylabel('Signature', fontsize=11)
    ax4.set_title('Variance Difference: Weighted - Unweighted\n(Shows where individual differences vary most)', 
                 fontsize=12, fontweight='bold')
    plt.colorbar(im4, ax=ax4, label='Variance Difference')
    
    # 5. Sample signature trajectories (mean ± std across patients)
    ax5 = plt.subplot(2, 3, 5)
    sample_sigs = [0, 5, 10, 15]
    for sig_idx in sample_sigs:
        if sig_idx < weighted_lambda_avg.shape[0]:
            unweighted_traj = unweighted_lambda_avg[sig_idx, :].numpy()
            weighted_traj = weighted_lambda_avg[sig_idx, :].numpy()
            unweighted_std = unweighted_lambda[:, sig_idx, :].std(dim=0).numpy()
            weighted_std = weighted_lambda[:, sig_idx, :].std(dim=0).numpy()
            
            ax5.plot(unweighted_traj, label=f'Sig {sig_idx} (Unweighted)', alpha=0.7, linewidth=1.5)
            ax5.fill_between(range(len(unweighted_traj)), 
                            unweighted_traj - unweighted_std, 
                            unweighted_traj + unweighted_std, alpha=0.2)
            ax5.plot(weighted_traj, label=f'Sig {sig_idx} (Weighted)', linestyle='--', alpha=0.7, linewidth=1.5)
            ax5.fill_between(range(len(weighted_traj)), 
                            weighted_traj - weighted_std, 
                            weighted_traj + weighted_std, alpha=0.2)
    
    ax5.set_xlabel('Time (Age)', fontsize=11)
    ax5.set_ylabel('Lambda Value (Mean ± SD)', fontsize=11)
    ax5.set_title('Lambda Trajectories: Sample Signatures\n(Mean ± SD across patients)', 
                 fontsize=12, fontweight='bold')
    ax5.legend(fontsize=7, ncol=2)
    ax5.grid(True, alpha=0.3)
    
    # 6. Correlation by signature (shows which signatures differ most)
    ax6 = plt.subplot(2, 3, 6)
    sig_correlations = []
    for sig_idx in range(weighted_lambda.shape[1]):
        sig_weighted = weighted_lambda[:, sig_idx, :].numpy().flatten()
        sig_unweighted = unweighted_lambda[:, sig_idx, :].numpy().flatten()
        sig_corr = np.corrcoef(sig_weighted, sig_unweighted)[0, 1]
        sig_correlations.append(sig_corr)
    
    ax6.bar(range(len(sig_correlations)), sig_correlations, alpha=0.7, edgecolor='black')
    ax6.axhline(individual_correlation, color='r', linestyle='--', linewidth=2, 
                label=f'Overall: {individual_correlation:.4f}')
    ax6.set_xlabel('Signature', fontsize=11)
    ax6.set_ylabel('Correlation', fontsize=11)
    ax6.set_title('Lambda Correlation by Signature\n(Individual values per signature)', 
                 fontsize=12, fontweight='bold')
    ax6.legend()
    ax6.grid(True, alpha=0.3, axis='y')
    
    plt.tight_layout()
    plt.show()
    
    print(f"\n✅ Individual lambda comparison complete")
    print(f"   Overall correlation (N×K×T): {individual_correlation:.6f}")
    if individual_correlation > 0.99:
        print(f"   ⚠️  High correlation observed, but individual differences exist")
        print("   (see variance difference heatmap and distribution)")
    else:
        print(f"   Lambda shows clear changes with IPW")
    
else:
    print("⚠️  Model files not found. Check paths:")
    print(f"  Weighted: {weighted_model_path}")
    print(f"  Unweighted: {unweighted_model_path}")

================================================================================
LAMBDA COMPARISON: Weighted vs Unweighted Models (Individual Level)
================================================================================

Loading weighted model: batch_00_model.pt
Loading unweighted model: enrollment_model_W0.0001_batch_0_10000.pt

Weighted lambda shape: torch.Size([10000, 21, 52])
Unweighted lambda shape: torch.Size([10000, 21, 52])
Using 10000 patients for comparison

Individual Lambda Comparison (N×K×T):
  Correlation: 0.987857
  Mean absolute difference: 0.128746
  Max absolute difference: 6.718026

✅ Individual lambda comparison complete
   Overall correlation (N×K×T): 0.987857
   Lambda shows clear changes with IPW

%run /Users/sarahurbut/aladynoulli2/pyScripts/dec_6_revision/new_notebooks/pythonscripts/create_S29_ipw_analysis.py

================================================================================
CREATING SUPPLEMENTARY FIGURE S29: IPW ANALYSIS
================================================================================

1. Loading and plotting IPW weights distribution...
   ✓ Loaded 469,553 weights
   Mean: 0.930, Std: 1.028
   Min: 0.169, Max: 6.631

2. Loading prevalences...
   ✓ Loaded prevalences

3. Loading models (weighted and unweighted)...
   ✓ Loaded and processed 10 batches

4. Creating combined S29 figure...

✓ Saved S29 figure to: /Users/sarahurbut/aladynoulli2/pyScripts/dec_6_revision/new_notebooks/results/paper_figs/supp/s29/S29.pdf

================================================================================
S29 COMPLETE
================================================================================
✓ IPW weights distribution
✓ Phi/Pi/Prevalence comparison (phi stable, pi/prevalence can change)
✓ Lambda comparison (6 panels showing individual differences)

This figure demonstrates the full IPW story:
  - Weights distribution shows the reweighting scheme
  - Phi remains stable (signature structure preserved)
  - Lambda/Pi adapts (model adjusts to reweighted population)
  - Prevalence changes (population demographics shift)

"""
Compare Phi, Pi, and Prevalence: Demonstrating IPW Effects

3-column plot showing:
1. Phi (weighted vs unweighted) - averaged over all signatures, same initialization → STABLE
2. Pi (weighted vs unweighted) - same initialization, but lambda adapts → CAN CHANGE
3. Prevalence (weighted vs unweighted) - all 400K patients → CAN CHANGE

This demonstrates that phi remains stable while lambda/pi and prevalence adapt to IPW.
"""

import torch
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from pathlib import Path
import sys

# Add path for utils
sys.path.append('/Users/sarahurbut/aladynoulli2/pyScripts')
from utils import calculate_pi_pred

print("="*80)
print("COMPARING PHI, PI, AND PREVALENCE: IPW EFFECTS")
print("="*80)

# Data directories
data_dir = Path('/Users/sarahurbut/Library/CloudStorage/Dropbox-Personal/data_for_running/')
model_1218_dir = Path('/Users/sarahurbut/Library/CloudStorage/Dropbox-Personal/batch_models_weighted_vec_censoredE_1218/')
unweighted_model_dir = Path('/Users/sarahurbut/Library/CloudStorage/Dropbox/censor_e_batchrun_vectorized/')

# Load prevalences (all 400K)
print("\n1. Loading prevalences (all 400K patients)...")
weighted_prevalence_path = data_dir / 'prevalence_t_weighted_corrected.pt'
unweighted_prevalence_path = data_dir / 'prevalence_t_corrected.pt'

if weighted_prevalence_path.exists():
    prevalence_t_weighted = torch.load(str(weighted_prevalence_path), weights_only=False)
    if torch.is_tensor(prevalence_t_weighted):
        prevalence_t_weighted = prevalence_t_weighted.numpy()
    print(f"   ✓ Loaded weighted prevalence: {prevalence_t_weighted.shape}")
else:
    raise FileNotFoundError(f"Weighted prevalence not found: {weighted_prevalence_path}")

if unweighted_prevalence_path.exists():
    prevalence_t_unweighted = torch.load(str(unweighted_prevalence_path), weights_only=False)
    if torch.is_tensor(prevalence_t_unweighted):
        prevalence_t_unweighted = prevalence_t_unweighted.numpy()
    print(f"   ✓ Loaded unweighted prevalence: {prevalence_t_unweighted.shape}")
else:
    raise FileNotFoundError(f"Unweighted prevalence not found: {unweighted_prevalence_path}")

# Load models to get phi and pi
print("\n2. Loading models (weighted and unweighted) to extract phi and pi...")
n_batches = 10

phi_1218_list = []
phi_unweighted_list = []
pi_1218_list = []
pi_unweighted_list = []

for batch_idx in range(n_batches):
    # Model 1218
    path_1218 = model_1218_dir / f"batch_{batch_idx:02d}_model.pt"
    if not path_1218.exists():
        path_1218 = model_1218_dir / f"batch_{batch_idx}_model.pt"
    
    # Unweighted model
    path_unweighted = unweighted_model_dir / f"enrollment_model_W0.0001_batch_{batch_idx*10000}_{(batch_idx+1)*10000}.pt"
    
    if path_1218.exists() and path_unweighted.exists():
        # Load 1218
        ckpt_1218 = torch.load(path_1218, weights_only=False, map_location='cpu')
        if 'model_state_dict' in ckpt_1218:
            phi_1218 = ckpt_1218['model_state_dict']['phi'].detach()
            lambda_1218 = ckpt_1218['model_state_dict']['lambda_'].detach()
            kappa_1218 = ckpt_1218['model_state_dict'].get('kappa', torch.tensor(1.0))
        else:
            phi_1218 = ckpt_1218['phi'].detach()
            lambda_1218 = ckpt_1218['lambda_'].detach()
            kappa_1218 = ckpt_1218.get('kappa', torch.tensor(1.0))
        
        if torch.is_tensor(kappa_1218):
            kappa_1218 = kappa_1218.item() if kappa_1218.numel() == 1 else kappa_1218.mean().item()
        
        # Load unweighted
        ckpt_unweighted = torch.load(path_unweighted, weights_only=False, map_location='cpu')
        if 'model_state_dict' in ckpt_unweighted:
            phi_unweighted = ckpt_unweighted['model_state_dict']['phi'].detach()
            lambda_unweighted = ckpt_unweighted['model_state_dict']['lambda_'].detach()
            kappa_unweighted = ckpt_unweighted['model_state_dict'].get('kappa', torch.tensor(1.0))
        else:
            phi_unweighted = ckpt_unweighted['phi'].detach()
            lambda_unweighted = ckpt_unweighted['lambda_'].detach()
            kappa_unweighted = ckpt_unweighted.get('kappa', torch.tensor(1.0))
        
        if torch.is_tensor(kappa_unweighted):
            kappa_unweighted = kappa_unweighted.item() if kappa_unweighted.numel() == 1 else kappa_unweighted.mean().item()
        
        # Store phi
        phi_1218_list.append(phi_1218)
        phi_unweighted_list.append(phi_unweighted)
        
        # Compute pi and average across patients
        pi_1218_batch = calculate_pi_pred(lambda_1218, phi_1218, kappa_1218)
        pi_unweighted_batch = calculate_pi_pred(lambda_unweighted, phi_unweighted, kappa_unweighted)
        
        pi_1218_avg = pi_1218_batch.mean(dim=0)  # [D, T]
        pi_unweighted_avg = pi_unweighted_batch.mean(dim=0)  # [D, T]
        
        pi_1218_list.append(pi_1218_avg)
        pi_unweighted_list.append(pi_unweighted_avg)
        
        if batch_idx == 0:
            print(f"   Batch {batch_idx}: phi shape {phi_1218.shape}, pi avg shape {pi_1218_avg.shape}")

# Average across batches
phi_1218_avg = torch.stack(phi_1218_list).mean(dim=0)  # [K, D, T]
phi_unweighted_avg = torch.stack(phi_unweighted_list).mean(dim=0)  # [K, D, T]

pi_1218_avg = torch.stack(pi_1218_list).mean(dim=0)  # [D, T]
pi_unweighted_avg = torch.stack(pi_unweighted_list).mean(dim=0)  # [D, T]

print(f"\n   ✓ Averaged across {len(phi_1218_list)} batches")
print(f"   Phi weighted: {phi_1218_avg.shape}")
print(f"   Phi unweighted: {phi_unweighted_avg.shape}")
print(f"   Pi weighted avg: {pi_1218_avg.shape}")
print(f"   Pi unweighted avg: {pi_unweighted_avg.shape}")

# Average phi across ALL signatures (like old notebook)
print("\n3. Averaging phi across all signatures for each disease...")
phi_1218_avg_over_sigs = phi_1218_avg.mean(dim=0)  # [D, T] - average over K dimension
phi_unweighted_avg_over_sigs = phi_unweighted_avg.mean(dim=0)  # [D, T]

print(f"   Phi weighted (avg over sigs): {phi_1218_avg_over_sigs.shape}")
print(f"   Phi unweighted (avg over sigs): {phi_unweighted_avg_over_sigs.shape}")

# Calculate overall correlations
print("\n4. Calculating overall correlations...")

# Phi correlation (averaged over all signatures)
phi_1218_flat = phi_1218_avg_over_sigs.numpy().flatten()
phi_unweighted_flat = phi_unweighted_avg_over_sigs.numpy().flatten()
phi_correlation = np.corrcoef(phi_1218_flat, phi_unweighted_flat)[0, 1]

# Pi correlation
pi_1218_flat = pi_1218_avg.numpy().flatten()
pi_unweighted_flat = pi_unweighted_avg.numpy().flatten()
pi_correlation = np.corrcoef(pi_1218_flat, pi_unweighted_flat)[0, 1]

# Prevalence correlation
prev_weighted_flat = prevalence_t_weighted.flatten()
prev_unweighted_flat = prevalence_t_unweighted.flatten()
valid_prev_mask = ~(np.isnan(prev_weighted_flat) | np.isnan(prev_unweighted_flat))
prev_correlation = np.corrcoef(prev_weighted_flat[valid_prev_mask], prev_unweighted_flat[valid_prev_mask])[0, 1]

print(f"   Phi correlation (weighted vs unweighted): {phi_correlation:.6f} (should be ~1.0, STABLE)")
print(f"   Pi correlation (weighted vs unweighted): {pi_correlation:.6f} (can differ, CAN CHANGE)")
print(f"   Prevalence correlation (weighted vs unweighted): {prev_correlation:.6f} (can differ, CAN CHANGE)")

# Plot comparison for selected diseases
DISEASES_TO_PLOT = [
    (112, "Myocardial Infarction"),
    (66, "Depression"),
    (16, "Breast cancer [female]"),
    (127, "Atrial fibrillation"),
    (47, "Type 2 diabetes"),
]

# Load disease names if available
disease_names_dict = {}
try:
    disease_names_path = Path("/Users/sarahurbut/aladynoulli2/pyScripts/dec_6_revision/new_notebooks/results/disease_names.csv")
    if disease_names_path.exists():
        disease_df = pd.read_csv(disease_names_path)
        disease_names_dict = dict(zip(disease_df['index'], disease_df['name']))
        print(f"✓ Loaded disease names")
except:
    pass

# Create figure: 3 columns (Phi | Pi | Prevalence), 5 rows (one per disease)
fig, axes = plt.subplots(len(DISEASES_TO_PLOT), 3, figsize=(18, 4*len(DISEASES_TO_PLOT)))
if len(DISEASES_TO_PLOT) == 1:
    axes = axes.reshape(1, -1)

time_points = np.arange(phi_1218_avg_over_sigs.shape[1]) + 30

for idx, (disease_idx, disease_name) in enumerate(DISEASES_TO_PLOT):
    if disease_idx >= phi_1218_avg_over_sigs.shape[0]:
        continue
    
    # Get disease name
    if disease_names_dict and disease_idx in disease_names_dict:
        display_name = disease_names_dict[disease_idx]
    else:
        display_name = disease_name
    
    # ===== COLUMN 1: Phi Comparison (averaged over all signatures) =====
    ax1 = axes[idx, 0]
    
    phi_1218_traj = phi_1218_avg_over_sigs[disease_idx, :].numpy()
    phi_unweighted_traj = phi_unweighted_avg_over_sigs[disease_idx, :].numpy()
    
    ax1.plot(time_points, phi_unweighted_traj, label='Unweighted Phi', 
            linewidth=2, alpha=0.8, color='blue')
    ax1.plot(time_points, phi_1218_traj, label='1218 Phi (Weighted training, Same init)', 
            linewidth=2, alpha=0.8, linestyle='--', color='red')
    
    ax1.set_xlabel('Age', fontsize=11)
    ax1.set_ylabel('Average Phi (across all signatures)', fontsize=11)
    ax1.set_title(f'{display_name}\nPhi: 1218 vs Unweighted (Same Init)', 
                 fontsize=12, fontweight='bold')
    ax1.legend(fontsize=9)
    ax1.grid(True, alpha=0.3)
    
    # Add correlation annotation
    disease_phi_corr = np.corrcoef(phi_unweighted_traj, phi_1218_traj)[0, 1]
    disease_phi_diff = np.abs(phi_1218_traj - phi_unweighted_traj).mean()
    ax1.text(0.02, 0.98, f'Corr: {disease_phi_corr:.4f}\nMean diff: {disease_phi_diff:.4f}', 
            transform=ax1.transAxes, verticalalignment='top', fontsize=9,
            bbox=dict(boxstyle='round', facecolor='lightgreen', alpha=0.8))
    
    # ===== COLUMN 2: Pi Comparison =====
    ax2 = axes[idx, 1]
    
    pi_1218_traj = pi_1218_avg[disease_idx, :].numpy()
    pi_unweighted_traj = pi_unweighted_avg[disease_idx, :].numpy()
    
    ax2.plot(time_points, pi_unweighted_traj, label='Unweighted Pi', 
            linewidth=2, alpha=0.8, color='blue')
    ax2.plot(time_points, pi_1218_traj, label='1218 Pi (Weighted training, Same init)', 
            linewidth=2, alpha=0.8, linestyle='--', color='red')
    
    ax2.set_xlabel('Age', fontsize=11)
    ax2.set_ylabel('Average Pi (Disease Hazard)', fontsize=11)
    ax2.set_title(f'{display_name}\nPi: 1218 vs Unweighted (Same Init)', 
                 fontsize=12, fontweight='bold')
    ax2.legend(fontsize=9)
    ax2.grid(True, alpha=0.3)
    ax2.set_yscale('log')
    
    # Add correlation annotation
    disease_pi_corr = np.corrcoef(pi_unweighted_traj, pi_1218_traj)[0, 1]
    disease_pi_diff = np.abs(pi_1218_traj - pi_unweighted_traj).mean()
    ax2.text(0.02, 0.98, f'Corr: {disease_pi_corr:.4f}\nMean diff: {disease_pi_diff:.4f}', 
            transform=ax2.transAxes, verticalalignment='top', fontsize=9,
            bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8))
    
    # ===== COLUMN 3: Prevalence Comparison =====
    ax3 = axes[idx, 2]
    
    if disease_idx < prevalence_t_weighted.shape[0] and disease_idx < prevalence_t_unweighted.shape[0]:
        weighted_prev_traj = prevalence_t_weighted[disease_idx, :]
        unweighted_prev_traj = prevalence_t_unweighted[disease_idx, :]
        
        # Match time points (prevalence might have different T)
        min_T = min(len(weighted_prev_traj), len(unweighted_prev_traj), len(time_points))
        time_points_prev = time_points[:min_T]
        weighted_prev_traj = weighted_prev_traj[:min_T]
        unweighted_prev_traj = unweighted_prev_traj[:min_T]
        
        ax3.plot(time_points_prev, unweighted_prev_traj, label='Unweighted Prevalence', 
                linewidth=2, alpha=0.8, color='blue')
        ax3.plot(time_points_prev, weighted_prev_traj, label='Weighted Prevalence (IPW)', 
                linewidth=2, alpha=0.8, linestyle='--', color='red')
        
        ax3.set_xlabel('Age', fontsize=11)
        ax3.set_ylabel('Prevalence', fontsize=11)
        ax3.set_title(f'{display_name}\nPrevalence: Weighted vs Unweighted (All 400K)', 
                     fontsize=12, fontweight='bold')
        ax3.legend(fontsize=9)
        ax3.grid(True, alpha=0.3)
        ax3.set_yscale('log')
        
        # Add correlation annotation
        valid_mask = ~(np.isnan(weighted_prev_traj) | np.isnan(unweighted_prev_traj))
        if valid_mask.sum() > 0:
            disease_prev_corr = np.corrcoef(unweighted_prev_traj[valid_mask], weighted_prev_traj[valid_mask])[0, 1]
            disease_prev_diff = np.abs(weighted_prev_traj[valid_mask] - unweighted_prev_traj[valid_mask]).mean()
            ax3.text(0.02, 0.98, f'Corr: {disease_prev_corr:.4f}\nMean diff: {disease_prev_diff:.4f}', 
                    transform=ax3.transAxes, verticalalignment='top', fontsize=9,
                    bbox=dict(boxstyle='round', facecolor='lightblue', alpha=0.8))
    else:
        ax3.text(0.5, 0.5, f'Disease {disease_idx}\nnot found', 
               transform=ax3.transAxes, ha='center', va='center', fontsize=12)
        ax3.set_title(f'{display_name}\nPrevalence: Weighted vs Unweighted', 
                     fontsize=12, fontweight='bold')

plt.suptitle(f'Phi, Pi, and Prevalence: IPW Effects\n'
            f'Phi Correlation: {phi_correlation:.4f} (STABLE - Same Init) | '
            f'Pi Correlation: {pi_correlation:.4f} (CAN CHANGE - Lambda Adapts) | '
            f'Prevalence Correlation: {prev_correlation:.4f} (CAN CHANGE - Population Demographics)', 
            fontsize=14, fontweight='bold')
plt.tight_layout()

# Save plot
output_dir = Path('/Users/sarahurbut/aladynoulli2/pyScripts/dec_6_revision/new_notebooks/paper_figs/supp/s29/')
output_dir.mkdir(parents=True, exist_ok=True)
plot_path = output_dir / 'phi_pi_prevalence_ipw_effects.pdf'
plt.savefig(plot_path, dpi=300, bbox_inches='tight')
print(f"\n✓ Saved comparison plot to: {plot_path}")
plt.show()

print(f"\n{'='*80}")
print("SUMMARY")
print("="*80)
print(f"✓ Phi correlation (1218 vs unweighted): {phi_correlation:.6f} - STABLE (same initialization)")
print(f"✓ Pi correlation (1218 vs unweighted): {pi_correlation:.6f} - CAN CHANGE (lambda adapts with IPW)")
print(f"✓ Prevalence correlation (weighted vs unweighted): {prev_correlation:.6f} - CAN CHANGE (population demographics)")
print(f"\nKey Insight:")
print(f"  - Phi remains stable when initialized the same (signature structure preserved)")
print(f"  - Pi changes because lambda adapts to IPW reweighting (model adjusts to population)")
print(f"  - Prevalence changes because IPW shifts population demographics")

================================================================================
COMPARING PHI, PI, AND PREVALENCE: IPW EFFECTS
================================================================================

1. Loading prevalences (all 400K patients)...
   ✓ Loaded weighted prevalence: (348, 52)
   ✓ Loaded unweighted prevalence: (348, 52)

2. Loading models (weighted and unweighted) to extract phi and pi...
   Batch 0: phi shape torch.Size([21, 348, 52]), pi avg shape torch.Size([348, 52])

   ✓ Averaged across 10 batches
   Phi weighted: torch.Size([21, 348, 52])
   Phi unweighted: torch.Size([21, 348, 52])
   Pi weighted avg: torch.Size([348, 52])
   Pi unweighted avg: torch.Size([348, 52])

3. Averaging phi across all signatures for each disease...
   Phi weighted (avg over sigs): torch.Size([348, 52])
   Phi unweighted (avg over sigs): torch.Size([348, 52])

4. Calculating overall correlations...
   Phi correlation (weighted vs unweighted): 0.999948 (should be ~1.0, STABLE)
   Pi correlation (weighted vs unweighted): 0.995792 (can differ, CAN CHANGE)
   Prevalence correlation (weighted vs unweighted): 0.996047 (can differ, CAN CHANGE)

✓ Saved comparison plot to: /Users/sarahurbut/aladynoulli2/pyScripts/dec_6_revision/new_notebooks/paper_figs/supp/s29/phi_pi_prevalence_ipw_effects.pdf

================================================================================
SUMMARY
================================================================================
✓ Phi correlation (1218 vs unweighted): 0.999948 - STABLE (same initialization)
✓ Pi correlation (1218 vs unweighted): 0.995792 - CAN CHANGE (lambda adapts with IPW)
✓ Prevalence correlation (weighted vs unweighted): 0.996047 - CAN CHANGE (population demographics)

Key Insight:
  - Phi remains stable when initialized the same (signature structure preserved)
  - Pi changes because lambda adapts to IPW reweighting (model adjusts to population)
  - Prevalence changes because IPW shifts population demographics

import numpy as np
import matplotlib.pyplot as plt
from pathlib import Path
from IPython.display import Image, display, HTML
import subprocess

# Load IPW correction demonstration results
results_dir = Path('/Users/sarahurbut/aladynoulli2/pyScripts/dec_6_revision/new_notebooks/results')

print("="*80)
print("IPW RECOVERY DEMONSTRATION: Dropping 90% of Women")
print("="*80)
print("\nThis demonstration shows that IPW reweighting can recover correct patterns")
print("from artificially biased populations. We dropped 90% of women and showed")
print("that IPW recovers both prevalence and model parameters.\n")

# Helper function to display PDF (convert to PNG if needed)
def display_pdf_as_image(pdf_path):
    """Display PDF by converting to PNG or using HTML embed."""
    pdf_path = Path(pdf_path)
    png_path = pdf_path.with_suffix('.png')
    
    # Try to convert PDF to PNG if PNG doesn't exist
    if not png_path.exists() and pdf_path.exists():
        try:
            subprocess.run(['pdftoppm', '-png', '-singlefile', str(pdf_path), str(png_path.with_suffix(''))], 
                         check=True, capture_output=True)
            png_path = Path(str(png_path.with_suffix('')) + '.png')
        except (subprocess.CalledProcessError, FileNotFoundError):
            # If conversion fails, use HTML embed
            print(f"⚠️  Could not convert PDF to PNG. Displaying as HTML link:")
            display(HTML(f'<a href="{pdf_path}" target="_blank">Open {pdf_path.name}</a>'))
            return
    
    # Display PNG if it exists
    if png_path.exists():
        display(Image(filename=str(png_path)))
    elif pdf_path.exists():
        # Fallback: HTML link
        display(HTML(f'<a href="{pdf_path}" target="_blank">Open {pdf_path.name}</a>'))
    else:
        print(f"⚠️  File not found: {pdf_path}")

# 1. Show prevalence recovery plot
print("\n" + "-"*80)
print("1. PREVALENCE RECOVERY")
print("-"*80)
prevalence_plot_path = results_dir / 'ipw_correction_demonstration.png'
if prevalence_plot_path.exists():
    print("\nThis plot shows prevalence trajectories for:")
    print("  • Full Population (Baseline) - black solid line")
    print("  • 90% Women Dropped, No Adjustment - red dashed line (drops substantially)")
    print("  • 90% Women Dropped, With IPW Reweighting - blue dotted line (recovers)\n")
    
    display_pdf_as_image(prevalence_plot_path)
    
    print("\n✅ Key Finding: IPW reweighting recovers prevalence to within ~108-130% of baseline")
    print("   The plots show clear recovery for diseases like Myocardial Infarction and Prostate cancer")
else:
    print(f"⚠️  Prevalence recovery plot not found: {prevalence_plot_path}")
    print("   Run demonstrate_ipw_correction.py to generate this plot")

================================================================================
IPW RECOVERY DEMONSTRATION: Dropping 90% of Women
================================================================================

This demonstration shows that IPW reweighting can recover correct patterns
from artificially biased populations. We dropped 90% of women and showed
that IPW recovers both prevalence and model parameters.


--------------------------------------------------------------------------------
1. PREVALENCE RECOVERY
--------------------------------------------------------------------------------

This plot shows prevalence trajectories for:
  • Full Population (Baseline) - black solid line
  • 90% Women Dropped, No Adjustment - red dashed line (drops substantially)
  • 90% Women Dropped, With IPW Reweighting - blue dotted line (recovers)

✅ Key Finding: IPW reweighting recovers prevalence to within ~108-130% of baseline
   The plots show clear recovery for diseases like Myocardial Infarction and Prostate cancer

import numpy as np
import matplotlib.pyplot as plt
from pathlib import Path
import pandas as pd
import torch
import sys
from scipy.ndimage import gaussian_filter1d

# Add path for utils
sys.path.append('/Users/sarahurbut/aladynoulli2/pyScripts')
from utils import calculate_pi_pred

# Load IPW correction demonstration results
results_dir = Path('/Users/sarahurbut/aladynoulli2/pyScripts/dec_6_revision/new_notebooks/results')
data_dir = Path('/Users/sarahurbut/Library/CloudStorage/Dropbox-Personal/data_for_running/')

print("="*80)
print("IPW RECOVERY DEMONSTRATION: Dropping 90% of Women")
print("="*80)
print("\nLoading phi/pi from batches 1-5, prevalence from full 400K...")

# Load disease names if available
disease_names_dict = {}
try:
    disease_names_path = Path("/Users/sarahurbut/aladynoulli2/pyScripts/dec_6_revision/new_notebooks/results/disease_names.csv")
    if disease_names_path.exists():
        disease_df = pd.read_csv(disease_names_path)
        if 'index' in disease_df.columns and 'name' in disease_df.columns:
            disease_names_dict = dict(zip(disease_df['index'], disease_df['name']))
        print(f"✓ Loaded {len(disease_names_dict)} disease names")
except:
    pass

# Define diseases to plot
DISEASES_TO_PLOT = [
    (21, "Prostate cancer [male]"),
    (112, "Myocardial Infarction"),
    (16, "Breast cancer")
]

# Load and average phi/pi across batches 1-5
print("\n1. Loading phi and pi from batches 1-5...")
phi_full_list = []
phi_biased_list = []
phi_biased_ipw_list = []
pi_full_list = []
pi_biased_list = []
pi_biased_ipw_list = []

N_train_per_batch = 20000
for batch_idx in range(1, 6):  # batches 1-5
    batch_dir = results_dir / f'batch_{batch_idx}'
    
    if batch_dir.exists():
        phi_full_list.append(np.load(batch_dir / 'phi_full.npy'))
        phi_biased_list.append(np.load(batch_dir / 'phi_biased.npy'))
        phi_biased_ipw_list.append(np.load(batch_dir / 'phi_biased_ipw.npy'))
        pi_full_list.append(np.load(batch_dir / 'pi_full.npy'))
        pi_biased_list.append(np.load(batch_dir / 'pi_biased.npy'))
        pi_biased_ipw_list.append(np.load(batch_dir / 'pi_biased_ipw.npy'))

# Average across batches
phi_full = np.mean(phi_full_list, axis=0)  # [K, D, T]
phi_biased = np.mean(phi_biased_list, axis=0)
phi_biased_ipw = np.mean(phi_biased_ipw_list, axis=0)

pi_full = np.mean(pi_full_list, axis=0)  # [D, T]
pi_biased = np.mean(pi_biased_list, axis=0)
pi_biased_ipw = np.mean(pi_biased_ipw_list, axis=0)

phi_full_avg = phi_full.mean(axis=0)  # [D, T]
phi_biased_avg = phi_biased.mean(axis=0)
phi_biased_ipw_avg = phi_biased_ipw.mean(axis=0)

print(f"✓ Loaded and averaged {len(phi_full_list)} batches")

# Prevalence computation function (same as in demonstrate_ipw_correction.py)
def compute_smoothed_prevalence_at_risk(Y, E_corrected, weights=None, window_size=5, smooth_on_logit=True):
    """Compute smoothed prevalence with at-risk filtering."""
    N, D, T = Y.shape
    prevalence_t = np.zeros((D, T))
    
    is_weighted = weights is not None
    if weights is not None:
        weights_norm = weights / weights.sum() * N
    
    for d in range(D):
        for t in range(T):
            at_risk_mask = (E_corrected[:, d] >= t)
            
            if at_risk_mask.sum() == 0:
                prevalence_t[d, t] = 0.0
                continue
            
            Y_at_risk = Y[at_risk_mask, d, t]
            
            if is_weighted:
                weights_at_risk = weights_norm[at_risk_mask]
                numerator = np.sum(weights_at_risk * Y_at_risk)
                denominator = np.sum(weights_at_risk)
                if denominator > 0:
                    prevalence_t[d, t] = numerator / denominator
                else:
                    prevalence_t[d, t] = 0.0
            else:
                prevalence_t[d, t] = Y_at_risk.mean()
    
    # Smooth
    for d in range(D):
        if smooth_on_logit:
            prev_d = prevalence_t[d, :]
            prev_d_clipped = np.clip(prev_d, 1e-6, 1 - 1e-6)
            logit_prev = np.log(prev_d_clipped / (1 - prev_d_clipped))
            logit_prev_smooth = gaussian_filter1d(logit_prev, sigma=window_size/3)
            prevalence_t[d, :] = 1 / (1 + np.exp(-logit_prev_smooth))
        else:
            prevalence_t[d, :] = gaussian_filter1d(prevalence_t[d, :], sigma=window_size/3)
    
    return prevalence_t

# Load full 400K data and recompute prevalence
print("\n2. Loading full 400K data and recomputing prevalence...")
n_patients = 400000
Y = torch.load(str(data_dir / 'Y_tensor.pt'), weights_only=False)
E_corrected = torch.load(str(data_dir / 'E_matrix_corrected.pt'), weights_only=False)

if torch.is_tensor(Y):
    Y = Y.numpy()
if torch.is_tensor(E_corrected):
    E_corrected = E_corrected.numpy()

Y = Y[:n_patients]
E_corrected = E_corrected[:n_patients]

# Load patient IDs and covariates
pids_csv_path = Path('/Users/sarahurbut/aladynoulli2/pyScripts/csv/processed_ids.csv')
pids_df = pd.read_csv(pids_csv_path)
pids = pids_df['eid'].values[:n_patients]

covariates_path = data_dir / 'baselinagefamh_withpcs.csv'
cov_df = pd.read_csv(covariates_path)
sex_col = 'sex'
cov_df = cov_df[['identifier', sex_col]].dropna(subset=['identifier'])
cov_df = cov_df.drop_duplicates(subset=['identifier'])
cov_map = cov_df.set_index('identifier')

# Identify women
is_female = np.zeros(n_patients, dtype=bool)
for i, pid in enumerate(pids):
    if pid in cov_map.index:
        sex_val = cov_map.at[pid, sex_col]
        if sex_val == 0 or sex_val == 'Female' or str(sex_val).lower() == 'female':
            is_female[i] = True

print(f"✓ Loaded full 400K data: {Y.shape}")
print(f"  Women: {is_female.sum():,} ({100*is_female.sum()/n_patients:.1f}%)")

# Full population prevalence (baseline)
print("  Computing full population prevalence...")
prevalence_full = compute_smoothed_prevalence_at_risk(
    Y, E_corrected, weights=None, window_size=5, smooth_on_logit=True
)

# Drop 90% of women (same logic as demonstrate_ipw_correction.py)
print("  Dropping 90% of women...")
np.random.seed(42)  # Same seed as in demonstrate script
female_indices = np.where(is_female)[0]
n_females_to_keep = int(len(female_indices) * 0.1)  # Keep only 10% = drop 90%
females_to_keep = np.random.choice(female_indices, size=n_females_to_keep, replace=False)
female_mask = np.zeros(n_patients, dtype=bool)
female_mask[females_to_keep] = True
male_mask = ~is_female

remaining_mask = male_mask | female_mask
Y_dropped = Y[remaining_mask]
E_dropped = E_corrected[remaining_mask]
is_female_dropped = is_female[remaining_mask]

print(f"  After drop: {remaining_mask.sum():,} patients ({is_female_dropped.sum():,} women)")

# Prevalence without IPW
print("  Computing prevalence without IPW...")
prevalence_biased = compute_smoothed_prevalence_at_risk(
    Y_dropped, E_dropped, weights=None, window_size=5, smooth_on_logit=True
)

# Compute IPW weights
n_women_full = is_female.sum()
n_men_full = (~is_female).sum()
n_women_dropped = is_female_dropped.sum()
n_men_dropped = (~is_female_dropped).sum()

prop_women_full = n_women_full / n_patients
prop_men_full = n_men_full / n_patients
prop_women_dropped = n_women_dropped / remaining_mask.sum()
prop_men_dropped = n_men_dropped / remaining_mask.sum()

ipw_weights = np.ones(remaining_mask.sum())
ipw_weights[is_female_dropped] = prop_women_full / (prop_women_dropped + 1e-10)
ipw_weights[~is_female_dropped] = prop_men_full / (prop_men_dropped + 1e-10)
ipw_weights = ipw_weights / ipw_weights.mean()

# Prevalence with IPW
print("  Computing prevalence with IPW...")
prevalence_biased_ipw = compute_smoothed_prevalence_at_risk(
    Y_dropped, E_dropped, weights=ipw_weights, window_size=5, smooth_on_logit=True
)

print(f"✓ Recomputed prevalence from full 400K dataset")

# Create 3-column plot
print("\n3. Creating 3-column plot (Phi, Pi, Prevalence)...")
time_points = np.arange(phi_full_avg.shape[1]) + 30

fig, axes = plt.subplots(len(DISEASES_TO_PLOT), 3, figsize=(18, 5*len(DISEASES_TO_PLOT)))
if len(DISEASES_TO_PLOT) == 1:
    axes = axes.reshape(1, -1)

for idx, (disease_idx, disease_name) in enumerate(DISEASES_TO_PLOT):
    if disease_idx >= phi_full_avg.shape[0]:
        continue
    
    display_name = disease_names_dict.get(disease_idx, disease_name) if disease_names_dict else disease_name
    
    # Column 1: Phi comparison
    ax1 = axes[idx, 0]
    phi_full_traj = phi_full_avg[disease_idx, :]
    phi_biased_traj = phi_biased_avg[disease_idx, :]
    phi_biased_ipw_traj = phi_biased_ipw_avg[disease_idx, :]
    
    ax1.plot(time_points, phi_full_traj, label='Full Population', 
            linewidth=2, color='black', linestyle='-')
    ax1.plot(time_points, phi_biased_traj, label='Biased (no IPW)', 
            linewidth=2, color='red', linestyle='--')
    ax1.plot(time_points, phi_biased_ipw_traj, label='Biased (with IPW)', 
            linewidth=2, color='blue', linestyle=':')
    ax1.set_xlabel('Age', fontsize=11)
    ax1.set_ylabel('Average Phi (across signatures)', fontsize=11)
    ax1.set_title(f'{display_name}\nPhi: Stable with Same Init', 
                 fontsize=12, fontweight='bold')
    ax1.legend(fontsize=9)
    ax1.grid(True, alpha=0.3)
    
    # Column 2: Pi comparison
    ax2 = axes[idx, 1]
    pi_full_traj = pi_full[disease_idx, :]
    pi_biased_traj = pi_biased[disease_idx, :]
    pi_biased_ipw_traj = pi_biased_ipw[disease_idx, :]
    
    ax2.plot(time_points, pi_full_traj, label='Full Population', 
            linewidth=2, color='black', linestyle='-')
    ax2.plot(time_points, pi_biased_traj, label='Biased (no IPW)', 
            linewidth=2, color='red', linestyle='--')
    ax2.plot(time_points, pi_biased_ipw_traj, label='Biased (with IPW)', 
            linewidth=2, color='blue', linestyle=':')
    ax2.set_xlabel('Age', fontsize=11)
    ax2.set_ylabel('Average Pi (Disease Hazard)', fontsize=11)
    ax2.set_title(f'{display_name}\nPi: IPW Recovers', 
                 fontsize=12, fontweight='bold')
    ax2.legend(fontsize=9)
    ax2.grid(True, alpha=0.3)
    ax2.set_yscale('log')
    
    # Column 3: Prevalence comparison (from full 400K)
    ax3 = axes[idx, 2]
    prev_full_traj = prevalence_full[disease_idx, :]
    prev_biased_traj = prevalence_biased[disease_idx, :]
    prev_biased_ipw_traj = prevalence_biased_ipw[disease_idx, :]
    
    ax3.plot(time_points, prev_full_traj, label='Full Population', 
            linewidth=2, color='black', linestyle='-')
    ax3.plot(time_points, prev_biased_traj, label='Biased (no IPW)', 
            linewidth=2, color='red', linestyle='--')
    ax3.plot(time_points, prev_biased_ipw_traj, label='Biased (with IPW)', 
            linewidth=2, color='blue', linestyle=':')
    ax3.set_xlabel('Age', fontsize=11)
    ax3.set_ylabel('Prevalence', fontsize=11)
    ax3.set_title(f'{display_name}\nPrevalence: IPW Recovers', 
                 fontsize=12, fontweight='bold')
    ax3.legend(fontsize=9)
    ax3.grid(True, alpha=0.3)
    ax3.set_yscale('log')

plt.suptitle('IPW Recovery: Full Population vs Biased Sample (with/without IPW)\n'
            'Phi/Pi: Pooled across 5 batches (20K each) | Prevalence: Full 400K dataset', 
            fontsize=14, fontweight='bold')
plt.tight_layout()

# Save plot
output_path = results_dir / 'ipw_recovery_phi_pi_prevalence.pdf'
plt.savefig(output_path, dpi=300, bbox_inches='tight')
print(f"\n✓ Saved plot to: {output_path}")
plt.show()

================================================================================
IPW RECOVERY DEMONSTRATION: Dropping 90% of Women
================================================================================

Loading phi/pi from batches 1-5, prevalence from full 400K...

1. Loading phi and pi from batches 1-5...
✓ Loaded and averaged 5 batches

2. Loading full 400K data and recomputing prevalence...
✓ Loaded full 400K data: (400000, 348, 52)
  Women: 217,458 (54.4%)
  Computing full population prevalence...
  Dropping 90% of women...
  After drop: 204,287 patients (21,745 women)
  Computing prevalence without IPW...
  Computing prevalence with IPW...
✓ Recomputed prevalence from full 400K dataset

3. Creating 3-column plot (Phi, Pi, Prevalence)...

✓ Saved plot to: /Users/sarahurbut/aladynoulli2/pyScripts/dec_6_revision/new_notebooks/results/ipw_recovery_phi_pi_prevalence.pdf

	Category	Unweighted	Weighted	Difference	Pct_Change
0	Age 60+	92.115054	83.974975	-8.140079	-8.836860
1	White British	89.330491	84.021691	-5.308800	-5.942876
2	University Degree	32.984562	28.661119	-4.323443	-13.107474
3	Good/Excellent Health	74.803909	72.866455	-1.937454	-2.590044

	Category	Unweighted	Weighted	Difference	Pct_Change
0	Age 60+	92.115054	83.974975	-8.140079	-8.836860
1	White British	89.330491	84.021691	-5.308800	-5.942876
2	University Degree	32.984562	28.661119	-4.323443	-13.107474
3	Good/Excellent Health	74.803909	72.866455	-1.937454	-2.590044

	Metric	Value
0	Mean Difference	0.003521
1	Std Difference	0.110672
2	Max Absolute Difference	1.431260
3	Mean Absolute Difference	0.086914

R1 Q1: Selection Bias / Socioeconomic Bias¶

Reviewer Question¶

Why This Matters¶

Our Approach¶

Key Findings¶

1. Inverse Probability Weighting Analysis¶

2. Impact on Model Signatures (Phi)¶

3. Impact on Population Prevalence Patterns (Lambda)¶

4. Impact on Disease Hazards (Pi)¶

5. IPW Recovery Demonstration: Artificially Induced Bias¶

3. Summary & Response Text¶

Key Findings¶

Response to Reviewer¶

References¶